"""
Pydantic schemas for GCR Solver Manager database models.
"""

from datetime import datetime
from typing import Optional, Dict, Any, List
from pydantic import BaseModel, Field, validator


class BaseSchema(BaseModel):
    """Base schema with common configuration."""
    
    class Config:
        from_attributes = True
        use_enum_values = True


# Build Configuration Schemas
class BuildConfigurationBase(BaseSchema):
    """Base build configuration schema."""
    
    name: str = Field(..., description="Configuration name")
    gpu_type: str = Field(..., description="GPU type (cuda/hip)")
    solver_type: str = Field(..., description="Solver type")
    debug_mode: bool = Field(default=False, description="Debug mode enabled")
    cmake_args: Dict[str, Any] = Field(default_factory=dict, description="CMake arguments")
    environment_vars: Dict[str, Any] = Field(default_factory=dict, description="Environment variables")
    source_hash: Optional[str] = Field(None, description="Source files hash")
    description: Optional[str] = Field(None, description="Configuration description")
    is_active: bool = Field(default=True, description="Configuration is active")
    
    @validator('gpu_type')
    def validate_gpu_type(cls, v):
        if v not in ['cuda', 'hip']:
            raise ValueError('GPU type must be either "cuda" or "hip"')
        return v
    
    @validator('solver_type')
    def validate_solver_type(cls, v):
        valid_types = ['gcr', 'ca-gcr', 'gmres', 'bca-gmres']
        if v not in valid_types:
            raise ValueError(f'Solver type must be one of: {valid_types}')
        return v


class BuildConfigurationCreate(BuildConfigurationBase):
    """Schema for creating build configurations."""
    pass


class BuildConfigurationUpdate(BaseSchema):
    """Schema for updating build configurations."""
    
    name: Optional[str] = None
    gpu_type: Optional[str] = None
    solver_type: Optional[str] = None
    debug_mode: Optional[bool] = None
    cmake_args: Optional[Dict[str, Any]] = None
    environment_vars: Optional[Dict[str, Any]] = None
    source_hash: Optional[str] = None
    description: Optional[str] = None
    is_active: Optional[bool] = None
    
    @validator('gpu_type')
    def validate_gpu_type(cls, v):
        if v is not None and v not in ['cuda', 'hip']:
            raise ValueError('GPU type must be either "cuda" or "hip"')
        return v
    
    @validator('solver_type')
    def validate_solver_type(cls, v):
        if v is not None:
            valid_types = ['gcr', 'ca-gcr', 'gmres', 'bca-gmres']
            if v not in valid_types:
                raise ValueError(f'Solver type must be one of: {valid_types}')
        return v


class BuildConfigurationResponse(BuildConfigurationBase):
    """Schema for build configuration responses."""
    
    id: int
    created_at: datetime
    updated_at: datetime


# Test Run Schemas
class TestRunBase(BaseSchema):
    """Base test run schema."""
    
    nproc: int = Field(..., description="Number of MPI processes")
    nprocx: Optional[int] = Field(None, description="nprocx parameter")
    resolution: float = Field(..., description="Resolution parameter")
    maxit: Optional[int] = Field(None, description="Maximum iterations")
    hostname: Optional[str] = Field(None, description="Hostname")
    user: Optional[str] = Field(None, description="Username")
    working_directory: Optional[str] = Field(None, description="Working directory")
    command_line: Optional[str] = Field(None, description="Command line")
    notes: Optional[str] = Field(None, description="Additional notes")


class TestRunCreate(TestRunBase):
    """Schema for creating test runs."""
    
    build_config_id: int = Field(..., description="Build configuration ID")


class TestRunUpdate(BaseSchema):
    """Schema for updating test runs."""
    
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    status: Optional[str] = None
    exit_code: Optional[int] = None
    slurm_job_id: Optional[str] = None
    node_info: Optional[Dict[str, Any]] = None
    duration_seconds: Optional[float] = None
    memory_usage_mb: Optional[float] = None
    cpu_usage_percent: Optional[float] = None
    error_message: Optional[str] = None
    notes: Optional[str] = None
    
    @validator('status')
    def validate_status(cls, v):
        if v is not None:
            valid_statuses = ['pending', 'running', 'completed', 'failed', 'cancelled']
            if v not in valid_statuses:
                raise ValueError(f'Status must be one of: {valid_statuses}')
        return v


class TestRunResponse(TestRunBase):
    """Schema for test run responses."""
    
    id: int
    build_config_id: int
    start_time: Optional[datetime]
    end_time: Optional[datetime]
    status: str
    exit_code: Optional[int]
    slurm_job_id: Optional[str]
    node_info: Dict[str, Any]
    duration_seconds: Optional[float]
    memory_usage_mb: Optional[float]
    cpu_usage_percent: Optional[float]
    error_message: Optional[str]
    created_at: datetime
    updated_at: datetime


# Log Entry Schemas
class LogEntryBase(BaseSchema):
    """Base log entry schema."""
    
    log_type: str = Field(..., description="Log type")
    log_level: str = Field(..., description="Log level")
    source: str = Field(..., description="Log source")
    file_path: Optional[str] = Field(None, description="Log file path")
    file_size: Optional[int] = Field(None, description="File size in bytes")
    content_hash: Optional[str] = Field(None, description="Content hash")
    is_compressed: bool = Field(default=False, description="Is compressed")
    content: Optional[str] = Field(None, description="Log content")
    category: Optional[str] = Field(None, description="Log category")
    tags: List[str] = Field(default_factory=list, description="Log tags")
    
    @validator('log_type')
    def validate_log_type(cls, v):
        valid_types = ['build', 'test', 'debug', 'slurm', 'error', 'analysis']
        if v not in valid_types:
            raise ValueError(f'Log type must be one of: {valid_types}')
        return v
    
    @validator('log_level')
    def validate_log_level(cls, v):
        valid_levels = ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']
        if v not in valid_levels:
            raise ValueError(f'Log level must be one of: {valid_levels}')
        return v


class LogEntryCreate(LogEntryBase):
    """Schema for creating log entries."""
    
    test_run_id: int = Field(..., description="Test run ID")


class LogEntryResponse(LogEntryBase):
    """Schema for log entry responses."""
    
    id: int
    test_run_id: int
    is_processed: bool
    processing_error: Optional[str]
    created_at: datetime
    updated_at: datetime


# Norm Data Schemas
class NormDataBase(BaseSchema):
    """Base norm data schema."""
    
    solver_tag: str = Field(..., description="Solver tag")
    iteration: Optional[int] = Field(None, description="Iteration number")
    step: Optional[int] = Field(None, description="Step number")
    norm_type: str = Field(..., description="Norm type")
    value: float = Field(..., description="Norm value")
    timestamp: Optional[datetime] = Field(None, description="Timestamp")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")


class NormDataCreate(NormDataBase):
    """Schema for creating norm data."""
    
    test_run_id: int = Field(..., description="Test run ID")
    log_entry_id: Optional[int] = Field(None, description="Log entry ID")


class NormDataResponse(NormDataBase):
    """Schema for norm data responses."""
    
    id: int
    test_run_id: int
    log_entry_id: Optional[int]
    created_at: datetime
    updated_at: datetime


# Performance Metric Schemas
class PerformanceMetricBase(BaseSchema):
    """Base performance metric schema."""
    
    metric_name: str = Field(..., description="Metric name")
    metric_type: str = Field(..., description="Metric type")
    value: float = Field(..., description="Metric value")
    unit: Optional[str] = Field(None, description="Metric unit")
    timestamp: Optional[datetime] = Field(None, description="Timestamp")
    phase: Optional[str] = Field(None, description="Execution phase")
    component: Optional[str] = Field(None, description="Component name")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")


class PerformanceMetricCreate(PerformanceMetricBase):
    """Schema for creating performance metrics."""
    
    test_run_id: int = Field(..., description="Test run ID")


class PerformanceMetricResponse(PerformanceMetricBase):
    """Schema for performance metric responses."""
    
    id: int
    test_run_id: int
    created_at: datetime
    updated_at: datetime


# SLURM Job Schemas
class SlurmJobBase(BaseSchema):
    """Base SLURM job schema."""
    
    job_id: str = Field(..., description="SLURM job ID")
    job_name: Optional[str] = Field(None, description="Job name")
    partition: Optional[str] = Field(None, description="Partition")
    account: Optional[str] = Field(None, description="Account")
    state: Optional[str] = Field(None, description="Job state")
    submit_time: Optional[datetime] = Field(None, description="Submit time")
    start_time: Optional[datetime] = Field(None, description="Start time")
    end_time: Optional[datetime] = Field(None, description="End time")
    nodes: Optional[str] = Field(None, description="Node list")
    cpus: Optional[int] = Field(None, description="CPU count")
    memory_mb: Optional[int] = Field(None, description="Memory in MB")
    time_limit_minutes: Optional[int] = Field(None, description="Time limit in minutes")
    exit_code: Optional[int] = Field(None, description="Exit code")
    exit_signal: Optional[int] = Field(None, description="Exit signal")
    stdout_path: Optional[str] = Field(None, description="Stdout file path")
    stderr_path: Optional[str] = Field(None, description="Stderr file path")
    script_path: Optional[str] = Field(None, description="Script file path")
    working_directory: Optional[str] = Field(None, description="Working directory")
    environment: Dict[str, Any] = Field(default_factory=dict, description="Environment variables")


class SlurmJobCreate(SlurmJobBase):
    """Schema for creating SLURM jobs."""
    
    test_run_id: int = Field(..., description="Test run ID")


class SlurmJobResponse(SlurmJobBase):
    """Schema for SLURM job responses."""
    
    id: int
    test_run_id: int
    created_at: datetime
    updated_at: datetime


# Build Cache Schemas
class BuildCacheBase(BaseSchema):
    """Base build cache schema."""
    
    cache_key: str = Field(..., description="Cache key")
    artifact_path: str = Field(..., description="Artifact path")
    build_time_seconds: Optional[float] = Field(None, description="Build time in seconds")
    artifact_size_bytes: Optional[int] = Field(None, description="Artifact size in bytes")
    is_valid: bool = Field(default=True, description="Cache is valid")
    hit_count: int = Field(default=0, description="Cache hit count")
    last_accessed: Optional[datetime] = Field(None, description="Last accessed time")


class BuildCacheCreate(BuildCacheBase):
    """Schema for creating build cache entries."""
    
    build_config_id: int = Field(..., description="Build configuration ID")


class BuildCacheResponse(BuildCacheBase):
    """Schema for build cache responses."""
    
    id: int
    build_config_id: int
    created_at: datetime
    updated_at: datetime


# Complex Response Schemas
class BuildConfigurationWithRuns(BuildConfigurationResponse):
    """Build configuration with associated test runs."""
    
    test_runs: List[TestRunResponse] = Field(default_factory=list)


class TestRunWithLogs(TestRunResponse):
    """Test run with associated log entries."""
    
    log_entries: List[LogEntryResponse] = Field(default_factory=list)
    norm_data: List[NormDataResponse] = Field(default_factory=list)
    performance_metrics: List[PerformanceMetricResponse] = Field(default_factory=list)


class TestRunSummary(BaseSchema):
    """Summary schema for test runs."""
    
    id: int
    build_config_name: str
    gpu_type: str
    solver_type: str
    debug_mode: bool
    nproc: int
    resolution: float
    status: str
    duration_seconds: Optional[float]
    created_at: datetime
    
    
# Query Schemas
class TestRunQuery(BaseSchema):
    """Schema for querying test runs."""
    
    build_config_id: Optional[int] = None
    gpu_type: Optional[str] = None
    solver_type: Optional[str] = None
    debug_mode: Optional[bool] = None
    status: Optional[str] = None
    start_date: Optional[datetime] = None
    end_date: Optional[datetime] = None
    limit: int = Field(default=100, ge=1, le=1000)
    offset: int = Field(default=0, ge=0)


class NormDataQuery(BaseSchema):
    """Schema for querying norm data."""
    
    test_run_id: Optional[int] = None
    solver_tag: Optional[str] = None
    norm_type: Optional[str] = None
    iteration: Optional[int] = None
    start_date: Optional[datetime] = None
    end_date: Optional[datetime] = None
    limit: int = Field(default=1000, ge=1, le=10000)
    offset: int = Field(default=0, ge=0)